package vn.gq.tym.utils;

import java.io.UnsupportedEncodingException;
import java.security.NoSuchAlgorithmException;

import javax.crypto.Cipher;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

public class AESCipher {

	private String iv = "fedcba9876543210";//Dummy iv (CHANGE IT!)
	private IvParameterSpec ivSpec;
	private SecretKeySpec keySpec;
	private Cipher cipher;
	public static final int BLOCK_SIZE = 32;
	
	public AESCipher(String key) {
		ivSpec = new IvParameterSpec(iv.getBytes());
		try {
			keySpec = new SecretKeySpec(padString(key).getBytes("cp1251"), "AES");
			cipher = Cipher.getInstance("AES/CBC/NoPadding");
		} catch (NoSuchAlgorithmException e) {
			e.printStackTrace();
		} catch (NoSuchPaddingException e) {
			e.printStackTrace();
		} catch (UnsupportedEncodingException e) {
			e.printStackTrace();
		}
	}
	
	public String encrypt(String plainText) throws Exception {
		return encrypt(plainText, Base64.DEFAULT);
	}
	
	public String encrypt(String plainText, int base64mode) throws Exception {
		if(plainText == null || plainText.length() == 0) throw new Exception("Empty string");
		cipher.init(Cipher.ENCRYPT_MODE, keySpec, ivSpec);
		byte[] enc = cipher.doFinal(padString(plainText).getBytes("cp1251"));
		byte[] b2enc = mergeByteArray(iv.getBytes(), enc);
		return Base64.encodeToString(b2enc, base64mode);
	}
	
	public String decrypt(String cipherText) throws Exception {
		return decrypt(cipherText, Base64.DEFAULT);
	}
	
	public String decrypt(String cipherText, int base64mode) throws Exception {
		if(cipherText == null || cipherText.length() == 0) throw new Exception("Empty string");
		String plainText = null;
		try {
			byte[] cipherByte = Base64.decode(cipherText, base64mode);
			byte[] ivByte = subByteArray(cipherByte, 0, 15);
			byte[] contentByte = subByteArray(cipherByte, 16, -1);
			ivSpec = new IvParameterSpec(ivByte);
			cipher.init(Cipher.DECRYPT_MODE, keySpec, ivSpec);
			plainText = unpadString(new String(cipher.doFinal(contentByte)));
		} catch (Exception e) {
			throw new Exception("[decrypt] " + e.getMessage());
		}
		return plainText;
	}
	
	public byte[] subByteArray(byte[] source, int firstPos, int lastPost) {
		byte[] subarr = null;
		if (source != null) {
			if (lastPost == -1) lastPost = source.length - 1;
			int size = lastPost - firstPos + 1;
			if (size > 0) {
				subarr = new byte[size];
				for(int i=0,j=i+firstPos;i<size;i++,j++) {
					subarr[i] = source[j];
				}
			}
		}
		return subarr;
	}
	
	public byte[] mergeByteArray(byte[] first, byte[] second) {
		if (first == null || second == null) return null;
		byte[] result = new byte[first.length + second.length];
		for (int i=0;i<first.length;i++) {
			result[i] = first[i];
		}
		for(int i=0, j=i+first.length;i<second.length;i++, j++) {
			result[j] = second[i];
		}
		return result;
	}
	
	private static String unpadString(String input) {
		int strLen = input.length() - (int) input.charAt(input.length() - 1);
		return input.substring(0, strLen);
	}

	private static String padString(String inputStr)
	{
	  int padLength = BLOCK_SIZE - inputStr.length() % BLOCK_SIZE;
	  char paddingChar = (char) padLength;
	  for (int i = 0; i < padLength; i++)
	  {
		  inputStr += paddingChar;
	  }
	  return inputStr;
	}
}
